import psycopg2


def add_column(table_name, conn):
    add_cell_sql = f"ALTER TABLE public.{table_name} add column if not exists _cellcode varchar(100)"
    add_comm_sql = f"ALTER TABLE public.{table_name} add column if not exists _commcode varchar(100)"
    add_street_sql = f"ALTER TABLE public.{table_name} add column if not exists _streetcode varchar(100)"
    add_dist_sql = f"ALTER TABLE public.{table_name} add column if not exists _distcode varchar(100)"
    with conn.cursor() as cursor:
        cursor.execute(add_cell_sql)
        print(f"图层 {table_name} 添加 _cellcode 字段成功, {add_cell_sql}")
        cursor.execute(add_comm_sql)
        print(f"图层 {table_name} 添加 _commcode 字段成功")
        cursor.execute(add_street_sql)
        print(f"图层 {table_name} 添加 _streetcode 字段成功")
        cursor.execute(add_dist_sql)
        print(f"图层 {table_name} 添加 _distcode 字段成功")


def update_cell_code_column(table_name, grid_name, grid_code, conn):
    update_cell_sql = f"UPDATE {table_name} jg SET _cellcode = grid.{grid_code} FROM {grid_name} grid WHERE ST_Intersects(jg.geom, grid.geom)"
    with conn.cursor() as cursor:
        cursor.execute(update_cell_sql)
        print(f"更新表 {table_name} 的 _cellcode 字段值成功")


def update_comm_code_column(table_name, grid_name, grid_code, conn):
    update_comm_sql = f"UPDATE {table_name} jg SET _commcode = grid.{grid_code} FROM {grid_name} grid WHERE ST_Intersects(jg.geom, grid.geom)"
    with conn.cursor() as cursor:
        cursor.execute(update_comm_sql)
        print(f"更新表 {table_name} 的 _commcode 字段值成功")


def update_street_code_column(table_name, grid_name, grid_code, conn):
    update_street_sql = f"UPDATE {table_name} jg SET _streetcode = grid.{grid_code} FROM {grid_name} grid WHERE ST_Intersects(jg.geom, grid.geom)"
    with conn.cursor() as cursor:
        cursor.execute(update_street_sql)


def update_dist_code_column(table_name, grid_name, grid_code, conn):
    update_dist_sql = f"UPDATE {table_name} jg SET _distcode = grid.{grid_code} FROM {grid_name} grid WHERE ST_Intersects(jg.geom, grid.geom)"
    with conn.cursor() as cursor:
        cursor.execute(update_dist_sql)
        print(f"更新表 {table_name} 的 _distcode 字段值成功")


def get_pg_tables():
    conn = psycopg2.connect(
        host="127.0.0.1",
        database="benxi",
        user="postgres",
        password="egova"
    )
    query_sql = """
            SELECT 
                f_table_name as table_name,
                type as geometry_type,
                srid as srid
            FROM geometry_columns
            WHERE f_table_schema NOT IN ('information_schema', 'pg_catalog')
        """
    names = []
    try:
        # 使用可滚动游标
        with conn.cursor() as cursor:
            cursor.execute(query_sql)
            while True:
                rows = cursor.fetchmany(100)
                if not rows:
                    break

                for row in rows:
                    table_name, type, srid = row
                    names.append((table_name, srid))

    except Exception as e:
        conn.rollback()
        print(f"查询失败: {e}")
    finally:
        conn.close()
    return names


def convert_srid(new_srid, conn):
    print(get_pg_tables())
    for name, srid in get_pg_tables():
        if (srid != new_srid):
            with conn.cursor() as cursor:
                cursor.execute(f"SELECT updategeometrysrid('{name}', 'geom', {new_srid})")
                print(f"指定{name}的srid为 {new_srid}")


def convert_multi(conn):
    convert_srid(4490, conn)
    for name, srid in get_pg_tables():
        if (name != cell_table_name or name != comm_table_name or name != street_table_name or name != dist_table_name):
            add_column(name, conn)
            update_cell_code_column(name, cell_table_name, cell_code_name, conn)
            update_comm_code_column(name, comm_table_name, comm_code_name, conn)
            update_street_code_column(name, street_table_name, street_code_name, conn)
            update_dist_code_column(name, dist_table_name, dist_code_name, conn)


def convert_one(table_name, cell_table_name, cell_code_name, comm_table_name, comm_code_name, street_table_name,
                street_code_name, dist_table_name, dist_code_name, conn):
    add_column(table_name, conn)
    update_cell_code_column(table_name, cell_table_name, cell_code_name, conn)
    update_comm_code_column(table_name, comm_table_name, comm_code_name, conn)
    update_street_code_column(table_name, street_table_name, street_code_name, conn)
    update_dist_code_column(table_name, dist_table_name, dist_code_name, conn)


if __name__ == '__main__':
    cell_table_name, cell_code_name = "单元网格", "bgcode"
    comm_table_name, comm_code_name = "社区", "sqcode"
    street_table_name, street_code_name = "街道", "jdcode"
    dist_table_name, dist_code_name = "区", "distid"

    conn = psycopg2.connect(
        host="127.0.0.1",
        database="benxi",
        user="postgres",
        password="egova"
    )
    # convert_one("路灯井盖", cell_table_name, cell_code_name, comm_table_name, comm_code_name, street_table_name,
    # street_code_name, dist_table_name, dist_code_name, conn)
    convert_multi(conn)
    conn.commit()
    conn.close()
